import io
import os
from flask import request
from PIL import Image
import imghdr
import numpy as np

from api import settings
from api.utils.api_utils import get_json_result

from deepdoc.vision.ocr import OCR

# 允许的图片拓展名
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}


# 判断文件是否为允许的图片格式
def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


def is_image(file):
    file.seek(0)
    file_type = imghdr.what(file)
    return file_type is not None


@manager.route('/resolve', methods=['POST'])  # noqa: F821
def ocr():
    if 'file' not in request.files:
        return get_json_result(data=False, message='No file part!', code=settings.RetCode.ARGUMENT_ERROR)
    file = request.files['file']

    if file.filename == '':
        return get_json_result(data=False, message='No file selected!', code=settings.RetCode.ARGUMENT_ERROR)

    if not allowed_file(file.filename) or not is_image(file):
        return get_json_result(data=False, message='File extension not allowed!', code=settings.RetCode.ARGUMENT_ERROR)

    ocr = OCR()
    img = Image.open(io.BytesIO(file.read())).convert('RGB')
    bxs = ocr(np.array(img))
    txt = "\n".join([t[0] for _, t in bxs if t[0]])
    return txt
